package cn.iocoder.yudao.module.infra.llm.impl.proxy;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.module.infra.api.file.FileApi;
import cn.iocoder.yudao.module.infra.llm.ILLMService;
import cn.iocoder.yudao.module.infra.llm.model.LLMMessage;
import cn.iocoder.yudao.module.infra.llm.model.LLMRes;
import cn.iocoder.yudao.module.resource.api.IMessageCollectionApi;
import cn.iocoder.yudao.module.resource.api.model.message.AudioDataDTO;
import cn.iocoder.yudao.module.resource.api.model.messagecollection.MessageCollectionDTO;
import cn.iocoder.yudao.module.resource.api.model.messagecollection.MessageCollectionResDTO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import jakarta.annotation.Resource;
import java.util.Arrays;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Service("llm-cache-proxy")
@Slf4j
public class LLMServiceCacheProxy implements ILLMService {
    @Resource
    private IMessageCollectionApi messageCollectionApi;

    @Resource
    private FileApi fileApi;

    private ILLMService targetService;

    public void setTargetService(ILLMService targetService) {
        this.targetService = targetService;
    }

    @Override
    public void streamChatResponse(Integer toneId, String prompt, String userText,
                                   List<LLMMessage> historyMessages,
                                   Consumer<LLMRes> onMessage) {
        final MessageCollectionResDTO messageRes;
        try {
            messageRes = messageCollectionApi.queryMessage(
                    0.7f, 20, userText, null
            );
        } catch (Exception e) {
            log.error(StrUtil.format("语意搜索异常, userText = {}", userText), e);
            targetService.streamChatResponse(toneId, prompt, userText, historyMessages, onMessage);
            return;
        }

        final List<MessageCollectionDTO> messages = messageRes.getMessages();
        if (CollUtil.isEmpty(messages)) {
            log.info("语意搜索未命中, userText = {}", userText);
            targetService.streamChatResponse(toneId, prompt, userText, historyMessages, onMessage);
            return;
        }

        log.info("语意搜索已命中, userText = {}, messageRes = {}", userText, JsonUtils.toJsonString(messageRes));
        final MessageCollectionDTO message = messages.get(RandomUtil.randomInt(0, messages.size()));
        final List<AudioDataDTO> audioDataList = message.getAudioData();
        if (CollUtil.isEmpty(audioDataList)) {
            log.info("语意搜索没音频, toneId = {}, userText = {}", toneId, userText);
            onMessage.accept(LLMRes.builder()
                    .fromLLMCache(true)
                    .done(true)
                    .completeSentence(true)
                    .cacheMessageId(message.getMessageId())
                    .replyText(message.getReplyText())
                    .cacheAudios(null)
                    .build());
            return;
        }

        final List<AudioDataDTO> toneAudios = audioDataList.stream()
                .filter(v -> StrUtil.isNotEmpty(v.getPath()) && v.getToneId() != null && v.getToneId().equals(toneId))
                .collect(Collectors.toList());
        if (CollUtil.isEmpty(toneAudios)) {
            log.info("语意搜索没音频, toneId = {}, userText = {}", toneId, userText);
            onMessage.accept(LLMRes.builder()
                    .fromLLMCache(true)
                    .done(true)
                    .completeSentence(true)
                    .cacheMessageId(message.getMessageId())
                    .replyText(message.getReplyText())
                    .cacheAudios(null)
                    .build());
            return;
        }

        log.info("语意搜索有音频, toneId = {}, userText = {}", toneId, userText);
        int index = 0;
        for (AudioDataDTO audio : toneAudios) {
            final boolean done = index == toneAudios.size() - 1;
            final String downloadUrl = fileApi.getTempDownloadUrl(audio.getEncryptKey(), audio.getPath(), null);
            audio.setUrl(downloadUrl);
            onMessage.accept(LLMRes.builder()
                    .fromLLMCache(true)
                    .fromTTSCache(true)
                    .done(done)
                    .completeSentence(true)
                    .cacheMessageId(message.getMessageId())
                    .replyText(audio.getSentence())
                    .cacheAudios(Arrays.asList(audio))
                    .build());
            index++;
        }
    }

}